#include "wavelet.h"
#include "cumacros.h"


//
// Convolve the rows and downsample the columns. Use half-point symmetrization.
//
inline static __device__ int
conv_row_down_col (data_t *y, data_t const *x, int m, int n, 
                   filt_t const* f, int lf)
{
	int ny = (n + lf-1) / 2;

	for (int i = threadIdx.x; i < m; i += blockDim.x){
		for (int j = 0; j < ny; ++j)
		{
			y[i + m*j] = 0.0f;

			for (int k = 0; k < lf; ++k)
			{
				int xj = 2*j+1 - (lf-1) + k;
				if (xj < 0) xj = -xj-1;
				if (xj >= n) xj = n-1 - (xj-n);

				y[i + m*j] += x[i + m*xj] * f[lf-1-k];
			}
		}
	}

	return ny;
}

//
// Convolve the columns and downsample the rows. Use half-point symmetrization.
//
inline static __device__ int
conv_col_down_row (data_t *a, data_t const *y, int m, int n, 
                   filt_t const* f, int lf)
{
	int ma = (m + lf-1) / 2;

	for (int i = threadIdx.x; i < ma; i += blockDim.x){
		for (int j = 0; j < n; ++j)
		{
			a[i + ma*j] = 0.0f;

			for (int k = 0; k < lf; ++k)
			{
				int yi =  2*i+1 - (lf-1) + k;
				if (yi < 0) yi = -yi-1;
				if (yi >= m) yi = m-1 - (yi-m);

				a[i + ma*j] += y[yi + m*j] * f[lf-1-k];
			}
		}
	}

	return ma;
}

//
// One level of forward wavelet decomposition
//
static inline __device__ void
forward_dwt2 (data_t *a, data_t *h, data_t *v, data_t *d,
              data_t const *x, data_t *y, int m, int n,
              filt_t const *lod, filt_t const *hid, int lf)
{

	int ny = conv_row_down_col (y, x, m, n, lod, lf);
	barrier ();

	conv_col_down_row (a, y, m, ny, lod, lf);
	barrier ();

	conv_col_down_row (h, y, m, ny, hid, lf);
	barrier ();

	conv_row_down_col (y, x, m, n, hid, lf);
	barrier ();

	conv_col_down_row (v, y, m, ny, lod, lf);
	barrier ();

	conv_col_down_row (d, y, m, ny, hid, lf);
	barrier ();
}

__device__ void
forward_transform (data_t *c, int nc, int *S,
                   data_t *x, int m, int n, int L,
                   filt_t const *lod, filt_t const *hid, int lf,
                   data_t *a, data_t *y, data_t *z)
{
	data_t *h, *v, *d, *tmp;

	for (int i = threadIdx.x; i < m*n; i += blockDim.x) z[i] = x[i];

	h = c + S[0]*S[1];
	for (int l = 1; l <= L; ++l){
		h += 3*S[0 + 2*l]*S[1 + 2*l];
	}

	barrier ();

	for (int l = L; l >= 1; --l)
	{
		int ml = S[0 + 2*l], nl = S[1 + 2*l];

		h = h - 3*ml*nl;
		v = h + ml*nl;
		d = v + ml*nl;

		forward_dwt2 (a, h, v, d, z, y, m, n, lod, hid, lf);

		tmp = z; z = a; a = tmp;
		m = ml; n = nl;

		barrier ();
	}


	for (int i = threadIdx.x; i < S[0]*S[1]; i += blockDim.x) c[i] = z[i];
}

__global__ void batch_forward_transform (
					data_t *c, int nc, int *S,
					data_t *x, int m, int n, int L,
					filt_t const *lod, filt_t const *hid, int lf,
					data_t *a, data_t *y, data_t *z)
{
	int b = blockIdx.x;

	forward_transform (c + b*nc, nc, S,
                       x + b*m*n, m, n,
                       L, lod, hid, lf,
                       a + b*m*n,
	                   y + b*m*n,
	                   z + b*m*n);
}

void
cu_forward_wavelet (data_t *C, int nc, data_t *Y, int m, int n,
                    int L, int *S, filt_t const *lod, filt_t const *hid, int lf,
                    int batch)
{
	data_t *a, *y, *z;
	filt_t *_lod, *_hid;

	cuda (Malloc (&a, 3*m*n*batch*sizeof(*a)));
	y = a + m*n*batch;
	z = y + m*n*batch;

	int *s;

	cuda (Malloc (&s, 2*(L+2)*sizeof(*S) + 2*lf*sizeof(*lod)));
	cuda (Memcpy (s, S, 2*(L+2)*sizeof(*S), cudaMemcpyHostToDevice));

	_lod = (filt_t*) (s+2*(L+2)); _hid = _lod+lf;

	cuda(Memcpy (_lod, lod, lf*sizeof(*lod), cudaMemcpyHostToDevice));
	cuda(Memcpy (_hid, hid, lf*sizeof(*hid), cudaMemcpyHostToDevice));

	batch_forward_transform <<< batch, 64 >>>
			(C, nc, s, Y, m, n, L, _lod, _hid, lf, a, y, z);
	cuda(GetLastError());

	cuda (Free (a));
	cuda (Free (s));
}

static inline __device__ int
conv_col_up_row (data_t *y, data_t const *a, int m, int n,
                 filt_t const *f, int lf)
{
	int my = 2*m-1 + lf-1;

	for (int i = threadIdx.x; i < my; i += blockDim.x){
		for (int j = 0; j < n; ++j)
		{
			for (int k = (i - (lf-1)) & 1; k < lf; k += 2){
				int ai = (i - (lf-1) + k) >> 1;
				if (ai >= 0 and ai < m)
					y[i + my*j] += a[ai + m*j] * f[lf-1-k];
			}
		}
	}

	return my;
}

static inline __device__ int
conv_row_up_col (data_t *y, data_t const *x, int m, int n,
                 filt_t const *f, int lf)
{
	int ny = 2*n-1 + lf-1;

	for (int i = threadIdx.x; i < m; i += blockDim.x){
		for (int j = 0; j < ny; ++j)
		{
			for (int k = (j - (lf-1)) & 1; k < lf; k += 2){
				int xj = (j - (lf-1) + k) >> 1;
				if (xj >= 0 and xj < n)
					y[i + m*j] += x[i + m*xj] * f[lf-1-k];
			}
		}
	}

	return ny;
}

static inline __device__ void
inverse_dwt2 (data_t *x,
              data_t *a, data_t *h, data_t *v, data_t *d,
              data_t *y, int m, int n,
              filt_t const *lor, filt_t const *hir, int lf)
{
	int mx = 2*m-1 + lf-1;
	for (int i = threadIdx.x; i < mx*n; i += blockDim.x) y[i] = 0.0f;

	int nx = 2*n-1 + lf-1;
	for (int i = threadIdx.x; i < mx*nx; i += blockDim.x) x[i] = 0.0f;

	barrier();

	conv_col_up_row (y, a, m, n, lor, lf);
	barrier();

	conv_col_up_row (y, h, m, n, hir, lf);
	barrier();

	conv_row_up_col (x, y, mx, n, lor, lf);
	barrier();

	for (int i = threadIdx.x; i < mx*n; i += blockDim.x) y[i] = 0.0f;
	barrier();

	conv_col_up_row (y, v, m, n, lor, lf);
	barrier();

	conv_col_up_row (y, d, m, n, hir, lf);
	barrier();

	conv_row_up_col (x, y, mx, n, hir, lf);
	barrier();
}

__device__ void
inverse_transform (data_t *x, int mx, int nx, 
                   int *S, int L, data_t *c, int nc,
                   filt_t const *lor, filt_t const *hir, int lf,
                   data_t *a, data_t *y, data_t *z)
{
	data_t *h, *v, *d;

	int m = S[0], n = S[1];
	for (int i = threadIdx.x; i < m*n; i += blockDim.x) a[i] = c[i];

	barrier();

	h = c + m*n;

	for (int l = 1; l < L+1; ++l)
	{
		m = S[0 + 2*l];
		n = S[1 + 2*l];

		v = h + m*n;
		d = v + m*n;

		inverse_dwt2 (z, a, h, v, d, y, m, n, lor, hir, lf);

		int om = (int) floor(((2*m-1 + lf-1) - S[0 + 2*(l+1)]) / 2.0f);
		int on = (int) floor(((2*n-1 + lf-1) - S[1 + 2*(l+1)]) / 2.0f);


		for (int i = threadIdx.x; i < S[0 + 2*(l+1)]; i += blockDim.x){
			for (int j = 0; j < S[1 + 2*(l+1)]; ++j){
				a[i + S[2*(l+1)]*j] = z[om+i + (2*m-1+lf-1)*(on+j)];
			}
		}

		h += 3*m*n;

		barrier();
	}

	for (int i = threadIdx.x; i < mx*nx; i += blockDim.x) x[i] = a[i]; 
}

__global__ void batch_inverse_transform
				(data_t *x, int m, int n, 
				 int *S, int L, data_t *c, int nc,
				 filt_t const *lor, filt_t const *hir, int lf,
				 data_t *a, data_t *y, data_t *z)

{
	int b = blockIdx.x;

	int ms = 2*((m + lf-1)/2)-1 + lf-1;
	int ns = 2*((n + lf-1)/2)-1 + lf-1;


	inverse_transform (x+b*m*n, m, n,
	                   S, L, c+b*nc, nc,
                       lor, hir, lf,
                       a+b*ms*ns,
	                   y+b*ms*ns,
	                   z+b*ms*ns);
}

void
cu_inverse_wavelet (data_t *X, int m, int n,
                    data_t *C, int nc, 
                    int L, int *S, filt_t const *lor, filt_t const *hir, int lf,
                    int batch)
{
	data_t *a, *y, *z;
	filt_t *_lor, *_hir;

	int ms = 2*((m + lf-1)/2)-1 + lf-1;
	int ns = 2*((n + lf-1)/2)-1 + lf-1;

	cuda (Malloc (&a, 3*ms*ns*batch*sizeof(*a)));
	y = a + ms*ns*batch;
	z = y + ms*ns*batch;

	int *s;

	cuda (Malloc (&s, 2*(L+2)*sizeof(*S) + 2*lf*sizeof(*lor)));
	cuda (Memcpy (s, S, 2*(L+2)*sizeof(*S), cudaMemcpyHostToDevice));

	_lor = (filt_t*) (s+2*(L+2)); _hir = _lor+lf;
	cuda(Memcpy (_lor, lor, lf*sizeof(*lor), cudaMemcpyHostToDevice));
	cuda(Memcpy (_hir, hir, lf*sizeof(*hir), cudaMemcpyHostToDevice));

	batch_inverse_transform <<< batch, 64 >>>
			(X, m, n, s, L, C, nc, _lor, _hir, lf, a, y, z);
	cuda(GetLastError());

	cuda (Free (a));
	cuda (Free (s));
}

